#!/usr/bin/env python

import os, sys
import numpy as np
from astropy.table import Table
import matplotlib.pylab as plt
from scipy import odr

def plotPRC(out_dir):
    tab = Table.read(f'{out_dir}/ptcdata.tab', format='ipac')

    model = odr.polynomial(1)

    fig = plt.figure(figsize=(20,40))

    for ch in range(16):
        row = ch // 4
        pos1 = row*8 + ch - row*4 + 1
        pos2 = pos1 + 4
        
    #    print(f'pos1: {pos1}, pos2: {pos2}')
        
        plt.subplot(8,4,pos1)
        
        index = tab['channel'] == ch+1
        
        t_exp = tab[index]['t_exp']
        flux = tab[index]['flux_4']
        
        # fit prc and calculate residuals
        data = odr.Data(t_exp, flux, we=1/(t_exp**2), wd=1/(flux**2) )
    #    data = odr.Data(t_exp, flux, we=1/(t_exp**2) )
    #    data = odr.Data(t_exp, flux)
        fitter = odr.ODR(data, model, ifixb=[1,1], beta0=[2500,800])
        res = fitter.run()
        poly = np.poly1d(res.beta[::-1])
        nonlin = (flux-poly(t_exp)) / poly(t_exp) * 100

        plt.scatter(t_exp, flux/1e3, zorder=1, alpha=0.85, color='m', edgecolors='none', s=30)
    #    plt.plot(t_exp, flux/1e3, 'g--o', alpha=0.75)
        plt.plot(t_exp, poly(t_exp)/1e3, 'b--', alpha=0.75)
        plt.xlabel('Expourse Time (s)')
        plt.ylabel('Flux (kDN)')
        plt.ylim([-1,6.1e1])
    #    plt.title('channel {}'.format(ch+1))
        plt.text(5,55,'channel {}'.format(ch+1), color='limegreen', fontsize=14)
        plt.grid()

        
        plt.subplot(8,4,pos2)
        print(res.beta)
        plt.scatter(t_exp, nonlin, alpha=0.75, color='darkorange', edgecolors='none', s=30)
        plt.hlines(y=1.8, xmin=t_exp.min(), xmax=t_exp.max(), linestyle=':', color='dimgrey')
        plt.hlines(y=-1.8, xmin=t_exp.min(), xmax=t_exp.max(), linestyle=':', color='dimgrey')
        plt.plot(t_exp, np.zeros(len(t_exp)), '--', color='grey')
        plt.ylim([-3.25,3.25])
        plt.xlabel('Expourse Time (s)')
        plt.ylabel('PRC Nonlinearity (%)')
    #    plt.title('channel {}'.format(ch+1))
        plt.text(5,2.5,'channel {}'.format(ch+1), color='brown', fontsize=14)
        plt.grid()
        
    plt.tight_layout()
    plt.savefig(f'{out_dir}/prc.pdf')
    plt.close()
    
if __name__ == '__main__':
    plotPRC(sys.argv[1])
